"""特征工程工具."""

from __future__ import annotations

import json
import logging
from typing import Any, Dict, Iterable, List, Optional

import numpy as np
import pandas as pd
from pydantic import BaseModel, Field

from ..features import FeatureSelector, build_feature_dataframe
from .etf_data import ETFDataTool
from .state import GLOBAL_DATA_STORE

logger = logging.getLogger(__name__)


def _format_feature_engineering_text(
    symbol: str,
    feature_rows: int,
    feature_columns: list[str],
    index_start: str | None,
    index_end: str | None,
    include_target: bool,
) -> str:
    """将特征工程结果格式化为易读的中文文本，便于理解."""
    lines = []
    
    # 标题
    lines.append(f"【{symbol} 特征工程完成报告】")
    lines.append("")
    
    # 基本信息
    lines.append("📊 特征数据概览：")
    lines.append(f"  • ETF 代码: {symbol}")
    lines.append(f"  • 特征样本数: {feature_rows} 条")
    lines.append(f"  • 特征维度: {len(feature_columns)} 个特征")
    if index_start and index_end:
        lines.append(f"  • 数据时间范围: {index_start} 至 {index_end}")
    lines.append("")
    
    # 目标标签信息
    if include_target:
        lines.append("🎯 目标标签：")
        lines.append("  • 已生成目标标签: target_up（下一交易日是否上涨）")
        lines.append("  • 说明: 1 表示上涨，0 表示下跌")
        lines.append("")
    
    # 特征分类统计
    lines.append("🔍 特征分类统计：")
    
    # 按特征类型分类统计
    feature_categories = {
        "移动平均": ["sma_", "ema_", "ma_cross", "ma_slope"],
        "技术指标": ["macd_", "rsi_", "atr_", "bb_", "kdj_", "cci_", "adx", "obv", "mfi_", "williams_r_"],
        "价格形态": ["price_range", "body_size", "shadow", "price_position"],
        "动量特征": ["mom_", "momentum_change", "acceleration"],
        "统计特征": ["skewness", "kurtosis", "quantile", "extreme"],
        "滞后特征": ["lag_", "_lag"],
        "时间特征": ["day_of_week", "month", "is_month_start", "is_month_end", "is_quarter_end"],
        "成交量特征": ["volume_", "turnover", "price_volume"],
        "波动率特征": ["volatility_", "vol_ratio", "realized_vol"],
        "交互特征": ["_interaction", "price_volume", "rsi_price", "macd_momentum"],
    }
    
    category_counts = {}
    for category, prefixes in feature_categories.items():
        count = sum(1 for col in feature_columns if any(col.startswith(prefix) or prefix in col for prefix in prefixes))
        if count > 0:
            category_counts[category] = count
    
    # 统计未分类的特征
    classified_cols = set()
    for prefixes in feature_categories.values():
        for col in feature_columns:
            if any(col.startswith(prefix) or prefix in col for prefix in prefixes):
                classified_cols.add(col)
    unclassified_count = len(feature_columns) - len(classified_cols)
    
    for category, count in sorted(category_counts.items()):
        lines.append(f"  • {category}: {count} 个特征")
    
    if unclassified_count > 0:
        lines.append(f"  • 其他特征: {unclassified_count} 个")
    
    lines.append("")
    
    # 主要特征列表（显示前20个）
    lines.append("📋 主要特征列表（前20个）：")
    display_cols = feature_columns[:20]
    for i, col in enumerate(display_cols, 1):
        lines.append(f"  {i:2d}. {col}")
    
    if len(feature_columns) > 20:
        lines.append(f"  ... 还有 {len(feature_columns) - 20} 个特征")
    
    lines.append("")
    
    # 使用提示
    lines.append("💡 使用提示：")
    lines.append(f"  • 特征数据已缓存，可以使用标准化代码 '{symbol}' 访问")
    lines.append("  • 可以使用 quant_analysis_tool 进行量化分析和模型训练")
    lines.append("  • 可以使用 prediction_tool 进行预测")
    if include_target:
        lines.append("  • 目标标签 target_up 可用于监督学习模型训练")
    
    return "\n".join(lines)


class FeatureEngineeringInput(BaseModel):
    """特征工程工具的输入模型."""

    symbol: str = Field(..., description="etf_data_tool 返回的标准化符号，例如 159919.SZ")
    additional_vol_windows: List[int] = Field(
        default_factory=list, description="额外波动率窗口"
    )
    include_target: bool = Field(
        default=False,
        description="是否生成下一交易日涨跌标签（仅使用 close 列）",
    )
    enable_feature_selection: bool = Field(
        default=True,
        description="是否启用特征选择，自动剔除低质量特征（默认 True）",
    )


class FeatureEngineering:
    """特征工程工具，用于基于缓存的 ETF 数据构建丰富的量化特征.
    
    这是一个普通的工具类，不依赖 LangChain 框架，可以直接使用。
    支持构建多种类型的特征，包括移动平均、技术指标、价格形态、动量特征等。
    """

    def __init__(self) -> None:
        """初始化特征工程工具."""
        pass

    def get_features(
        self,
        symbol: str,
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> Dict[str, Any]:
        """获取特征工程结果并返回字典格式.

        Args:
            symbol: ETF 标准化符号，例如 159919.SZ
            additional_vol_windows: 额外波动率窗口列表
            include_target: 是否生成下一交易日涨跌标签
            enable_feature_selection: 是否启用特征选择

        Returns:
            包含特征工程结果的字典

        Raises:
            RuntimeError: 如果数据未缓存
        """
        dataset = GLOBAL_DATA_STORE.get_dataset(symbol)
        if dataset is None:
            raise RuntimeError("请先调用 etf_data_tool 获取并缓存 ETF 历史数据")

        frames = build_feature_dataframe(
            dataset.frame,
            additional_windows=list(additional_vol_windows or []),
        )

        if include_target:
            # 使用中文列名（AkShare 数据格式）
            close = dataset.frame["收盘"].loc[frames.index]
            # 计算下一交易日是否上涨
            # shift(-1) 会将价格向前移动一位，最后一行会是 NaN（因为没有下一交易日）
            next_close = close.shift(-1)
            # 只有当 next_close 不是 NaN 时才进行比较，否则设为 NaN
            target_up = pd.Series(
                index=close.index,
                dtype="Int64",
            )
            # 只对有效的 next_close 进行比较
            valid_mask = pd.notna(next_close)
            target_up[valid_mask] = (next_close[valid_mask] > close[valid_mask]).astype("Int64")
            target_up[~valid_mask] = pd.NA
            frames["target_up"] = target_up
            # 删除包含 NaN 的行（最后一行）
            frames = frames.dropna(subset=["target_up"])

        # 特征选择（如果启用）
        feature_selection_result = None
        if enable_feature_selection:
            # 分离特征和目标变量
            if include_target and "target_up" in frames.columns:
                X = frames.drop(columns=["target_up"])
                y = frames["target_up"]
            else:
                X = frames
                y = None

            # 执行特征选择
            selector = FeatureSelector()
            feature_selection_result = selector.select_features(X, y)
            
            # 应用特征选择结果
            frames_selected = selector.apply_selection(frames, feature_selection_result)
            
            # 如果包含目标变量，需要重新添加
            if include_target and "target_up" in frames.columns:
                frames_selected["target_up"] = frames["target_up"]
            
            frames = frames_selected
            
            logger.info(
                f"特征选择完成: 原始特征数={feature_selection_result.original_count}, "
                f"选中特征数={feature_selection_result.selected_count}, "
                f"剔除特征数={len(feature_selection_result.removed_features)}"
            )

        GLOBAL_DATA_STORE.set_features(dataset.symbol, frames)
        
        # 准备输出数据
        payload = {
            "symbol": dataset.symbol,
            "feature_rows": len(frames),
            "feature_columns": list(frames.columns),
            "index_start": frames.index.min().isoformat() if not frames.empty else None,
            "index_end": frames.index.max().isoformat() if not frames.empty else None,
            "include_target": include_target,
        }
        
        # 添加特征选择信息
        if feature_selection_result:
            payload["feature_selection"] = {
                "enabled": True,
                "original_count": feature_selection_result.original_count,
                "selected_count": feature_selection_result.selected_count,
                "removed_count": len(feature_selection_result.removed_features),
                "removed_features": feature_selection_result.removed_features[:20],  # 只返回前20个
                "removal_reasons": {
                    feat: reason
                    for feat, reason in list(feature_selection_result.removal_reasons.items())[:20]
                },
            }
        else:
            payload["feature_selection"] = {"enabled": False}
        
        logger.info(
            "特征工程完成：symbol=%s, rows=%s, columns=%s",
            dataset.symbol,
            payload["feature_rows"],
            len(payload["feature_columns"]),
        )
        
        return payload

    def get_features_json(
        self,
        symbol: str,
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> str:
        """获取特征工程结果并返回 JSON 格式的字符串.

        Args:
            symbol: ETF 标准化符号，例如 159919.SZ
            additional_vol_windows: 额外波动率窗口列表
            include_target: 是否生成下一交易日涨跌标签
            enable_feature_selection: 是否启用特征选择

        Returns:
            JSON 格式的字符串
        """
        result = self.get_features(
            symbol=symbol,
            additional_vol_windows=additional_vol_windows,
            include_target=include_target,
            enable_feature_selection=enable_feature_selection,
        )
        return json.dumps(result, ensure_ascii=False)

    def get_features_text(
        self,
        symbol: str,
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> str:
        """获取特征工程结果并返回格式化的文本.

        Args:
            symbol: ETF 标准化符号，例如 159919.SZ
            additional_vol_windows: 额外波动率窗口列表
            include_target: 是否生成下一交易日涨跌标签
            enable_feature_selection: 是否启用特征选择

        Returns:
            格式化的中文文本
        """
        result = self.get_features(
            symbol=symbol,
            additional_vol_windows=additional_vol_windows,
            include_target=include_target,
            enable_feature_selection=enable_feature_selection,
        )
        return _format_feature_engineering_text(
            symbol=result["symbol"],
            feature_rows=result["feature_rows"],
            feature_columns=result["feature_columns"],
            index_start=result.get("index_start"),
            index_end=result.get("index_end"),
            include_target=result.get("include_target", False),
        )

    def run(
        self,
        symbol: str | Dict[str, Any],
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> str:
        """执行特征工程并返回 JSON 格式的字符串（向后兼容方法）.

        此方法支持两种调用方式：
        1. `run(symbol="159919.SZ", include_target=True)` - 直接传递参数
        2. `run({"symbol": "159919.SZ", "include_target": True})` - 传递字典（LangChain 兼容）

        Args:
            symbol: ETF 标准化符号或包含参数的字典
            additional_vol_windows: 额外波动率窗口列表（当 symbol 是字符串时使用）
            include_target: 是否生成下一交易日涨跌标签（当 symbol 是字符串时使用）
            enable_feature_selection: 是否启用特征选择（当 symbol 是字符串时使用）

        Returns:
            JSON 格式的字符串
        """
        # 支持字典参数（向后兼容 LangChain 调用方式）
        if isinstance(symbol, dict):
            symbol_str = symbol.get("symbol", "")
            additional_vol_windows_val = symbol.get("additional_vol_windows", None)
            include_target_val = symbol.get("include_target", False)
            enable_feature_selection_val = symbol.get("enable_feature_selection", True)
        else:
            symbol_str = symbol
            additional_vol_windows_val = additional_vol_windows
            include_target_val = include_target
            enable_feature_selection_val = enable_feature_selection

        return self.get_features_json(
            symbol=symbol_str,
            additional_vol_windows=additional_vol_windows_val,
            include_target=include_target_val,
            enable_feature_selection=enable_feature_selection_val,
        )

    # 向后兼容的方法别名
    def _get_output(
        self,
        symbol: str,
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> dict:
        """内部方法：获取原始输出字典（用于测试和调试）。"""
        return self.get_features(
            symbol=symbol,
            additional_vol_windows=additional_vol_windows,
            include_target=include_target,
            enable_feature_selection=enable_feature_selection,
        )

    # 向后兼容的方法别名（用于测试）
    def _run(
        self,
        symbol: str,
        additional_vol_windows: Iterable[int] | None = None,
        include_target: bool = False,
        enable_feature_selection: bool = True,
    ) -> str:
        """执行特征工程并返回格式化的文本（向后兼容方法，用于测试）。"""
        return self.get_features_text(
            symbol=symbol,
            additional_vol_windows=additional_vol_windows,
            include_target=include_target,
            enable_feature_selection=enable_feature_selection,
        )


# 向后兼容的类别名
FeatureEngineeringTool = FeatureEngineering

